// SPDX-License-Identifier: GPL-2.0+

#include <common.h>
#include <command.h>
#include <malloc.h>
#include <ext4fs.h>
#include <fs.h>

#include <jffs2/jffs2.h>

#include "loongson_update.h"

extern int mtdparts_init(void);
extern int ext4fs_find_file(const char *path, struct ext2fs_node *rootnode,
			struct ext2fs_node **foundnode, int expecttype);
extern int ext4fs_iterate_dir(struct ext2fs_node *dir, char *name,
			struct ext2fs_node **fnode, int *ftype);

DECLARE_GLOBAL_DATA_PTR;

const char *update_devname_str[UPDATE_DEV_COUNT] = {
	[UPDATE_DEV_USB]	= "usb",
	[UPDATE_DEV_TFTP]	= "tftp",
	[UPDATE_DEV_MMC]	= "mmc",
};

const char *update_typename_str[UPDATE_TYPE_COUNT] = {
	[UPDATE_TYPE_KERNEL]	= "kernel",
	[UPDATE_TYPE_ROOTFS]	= "rootfs",
	[UPDATE_TYPE_UBOOT]		= "uboot",
	[UPDATE_TYPE_DTB]		= "dtb",
	[UPDATE_TYPE_SYSTEM]	= "system",
	[UPDATE_TYPE_ALL]		= "all",
	[UPDATE_TYPE_BOOTSELECT]		= "bootselect",
	[UPDATE_TYPE_ARGS]		= "args",
	[UPDATE_TYPE_RESOLUTION]		= "resolution",
	[UPDATE_TYPE_ROTATION]		= "rotation",
};

static int mtdparts_probe(void)
{
	int ret = -1;

	run_command("sf probe", 0);

	ret = mtdparts_init();
	if (ret) {
		printf("-----> mtdparts init failed! partition not found or toolarge\r\n");
	}
	return ret;
}

static void user_env_save(void)
{
	char *env;
	char *mem;
	char *cmd;
#define ENV_BUF_LEN 1024

	printf("%s\r\n", __func__);

	env = env_get("syspart");//syspart 用于保存操作系统分区的分区号(linux 如sda1、sda2 ...) (uboot 如scsi 0:1、scsi 0:2 ...)
	if (env == NULL)
		return;

	mem = (char *)(memalign(ARCH_DMA_MINALIGN, ENV_BUF_LEN));
	cmd = (char *)(memalign(ARCH_DMA_MINALIGN, 128));

	sprintf(mem, "syspart=%s", env);

#ifdef LS_DOUBLE_SYSTEM
	env = env_get("syspart_last");
	sprintf(mem + strlen(mem) + 1, "syspart_last=%s", env ? env : env_get("syspart"));
#endif

	sprintf(cmd, "mtd erase spi_user"); //spi_user分区可用于存放用户自定义的数据，不会被擦除
	run_command(cmd, 0);
	sprintf(cmd, "mtd write spi_user 0x%p 0 0x%x", mem, ENV_BUF_LEN);
	run_command(cmd, 0);

	free(cmd);
	free(mem);
}

void user_env_fetch(void)
{
	char *env;
	char *mem;
	char *cmd;
	char *env_addr;
	char *env_user;	
	int ret;
#define ENV_BUF_LEN 1024

#ifdef LS_DOUBLE_SYSTEM
#else
	env = env_get("update"); //update 可用于判断环境变量是否被更新过，如spi flash的环境变量分区是否被擦除；
							 //也可用于判断uboot是否被更新过。update默认值是1，uboot第一次启动时可以把该变量修改为0，
							 //下次启动就不会再读取spi_user的用户分区
	if (strcmp(env, "0") == 0) {
		return;
	}
#endif

	printf("%s\r\n", __func__);

	env_user = (char *)(memalign(ARCH_DMA_MINALIGN, ENV_BUF_LEN));
	mem = (char *)(memalign(ARCH_DMA_MINALIGN, ENV_BUF_LEN));
	cmd = (char *)(memalign(ARCH_DMA_MINALIGN, 128));

	sprintf(cmd, "sf probe;mtd read spi_user 0x%p 0 0x%x", mem, ENV_BUF_LEN);
	run_command(cmd, 0);

	env_addr = (char *)gd->env_addr;
	gd->env_addr = (ulong)mem;
	ret = env_get_f("syspart", env_user, ENV_BUF_LEN);
	gd->env_addr = (ulong)env_addr;

#ifdef LS_DOUBLE_SYSTEM
	env = env_get("update"); //update 可用于判断环境变量是否被更新过，如spi flash的环境变量分区是否被擦除；
							 //也可用于判断uboot是否被更新过。update默认值是1，uboot第一次启动时可以把该变量修改为0，
							 //下次启动就不会再读取spi_user的用户分区
	if (strcmp(env, "0") != 0) {
		if (ret > 0) {
			env = env_get("syspart");
			if ((strcmp(env_user, env) != 0) || (env_get("syspart_last") == NULL)) {
				env_set("syspart", env_user);

				env_addr = (char *)gd->env_addr;
				gd->env_addr = (ulong)mem;
				ret = env_get_f("syspart_last", env_user, ENV_BUF_LEN);
				gd->env_addr = (ulong)env_addr;

				env_set("syspart_last", env_user ? env_user : env_get("syspart"));
			}
		}
		env_set("update", "0");
		env_save();
	} else {
		if (ret > 0) {
			env = env_get("syspart");
			if (strcmp(env_user, env) != 0) {
				env_set("syspart_last", env_user);
				env_save();
				user_env_save();
			}
		}
	}
#else
	if (ret > 0) {
		env = env_get("syspart");
		if (strcmp(env_user, env) != 0) {
			env_set("syspart", env_user);
		}
	}

	env_set("update", "0");
	env_save();
#endif

	free(cmd);
	free(mem);
	free(env_user);
}

static int update_uboot(int dev)
{
	int ret = -1;

	printf("update u-boot.............\n");

	switch (dev) {
	case UPDATE_DEV_USB:
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/u-boot.bin", 0);
		break;
	case UPDATE_DEV_TFTP:
		ret = run_command("tftpboot ${loadaddr} u-boot.bin", 0);
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/u-boot.bin", 0);
		break;
	}

	if (ret) {
		printf("upgrade uboot failed no /update/u-boot.bin file!\r\n");
		return ret;
	}

	if (mtdparts_probe()) {
		ret = run_command("sf erase 0 0xd8000;sf update ${loadaddr} 0 ${filesize}", 0);
	} else {
		char cmd[256];
		struct mtd_device *mtd_dev;
		struct part_info *part;
		u8 pnum;
		ulong uboot_size, env_size;

		find_dev_and_part("uboot", &mtd_dev, &pnum, &part);
		uboot_size = part->size;
		find_dev_and_part("uboot_env", &mtd_dev, &pnum, &part);
		env_size = part->size;

		sprintf(cmd, "sf erase uboot %lx;sf erase uboot_env %lx;sf update ${loadaddr} uboot ${filesize}", uboot_size, env_size);
		ret = run_command(cmd, 0);
	}

	user_env_save();

	if (ret) {
		printf("upgrade uboot failed!\r\n");
	} else {
		printf("upgrade uboot success\r\n");
	}

	return ret;
}

static int update_dtb(int dev)
{
	int ret = -1;

	printf("update dtb.............\n");

	switch (dev) {
	case UPDATE_DEV_USB:
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/dtb.bin", 0);
		break;
	case UPDATE_DEV_TFTP:
		ret = run_command("tftpboot ${loadaddr} dtb.bin", 0);
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/dtb.bin", 0);
		break;
	}

	if (ret) {
		printf("upgrade dtb failed no /update/dtb.bin file!\r\n");
		return ret;
	}

	if (mtdparts_probe()) {
		ret = -1;
	} else {
		ret = run_command("sf erase dtb 0x20000;sf update ${loadaddr} dtb ${filesize}", 0);
	}
	if (ret) {
		printf("upgrade dtb failed!\r\n");
	} else {
		printf("upgrade dtb success\r\n");
	}

	return ret;
}

static int update_kernel(int dev, char *typename, char *kernelname)
{
	char cmd[256];
	char kernel[32] = "uImage";
	int ret = -1;

	printf("update kernel....................\n");

	if (kernelname) {
		memset(kernel, 0, sizeof(kernel));
		sprintf(kernel, "%s", kernelname);
	}

	switch (dev) {
	case UPDATE_DEV_USB:
		sprintf(cmd, "fatload usb 0:1 ${loadaddr} update/%s", kernel);
		break;
	case UPDATE_DEV_TFTP:
		sprintf(cmd, "tftpboot ${loadaddr} %s", kernel);
		break;
	case UPDATE_DEV_MMC:
		sprintf(cmd, "fatload mmc 0:1 ${loadaddr} update/%s", kernel);
		break;
	}

	ret = run_command(cmd, 0);
	if (ret) {
		printf("upgrade kernel failed no update/%s file!\r\n", kernel);
		return ret;
	}

	if (strcmp(typename, "nand") == 0) {
		ret = mtdparts_probe();
		if (ret) {
			return ret;
		}
		printf("update kernel to nand............\n");
		ret = run_command("nand erase.part kernel;nand write ${loadaddr} kernel ${filesize}", 0);
		if (ret) {
			goto err;
		}
	} else if (strcmp(typename, "sata") == 0) {
		struct ext2fs_node *fdiro = NULL;
		int type;
		char buf[8];

		printf("update kernel to ssd.............\n");

		ret = run_command("scsi scan", 0);
		if (ret) {
			goto err;
		}

		sprintf(buf, "0:%s", env_get("syspart"));
		ret = fs_set_blk_dev("scsi", buf, FS_TYPE_EXT);
		if (ret) {
			goto err;
		}
		
		ext4fs_find_file("/", &ext4fs_root->diropen, &fdiro, FILETYPE_DIRECTORY);
		ext4fs_iterate_dir(fdiro, "boot", &fdiro, &type);
		ext4fs_free_node(fdiro, &ext4fs_root->diropen);

		if (type == FILETYPE_DIRECTORY) {
			sprintf(cmd, "ext4write scsi 0:${syspart} ${loadaddr} /boot/%s ${filesize}", kernel);
		} else {
			sprintf(cmd, "ext4write scsi 0:${syspart} ${loadaddr} /%s ${filesize}", kernel);
		}
		ret = run_command(cmd, 0);
		if (ret) {
			goto err;
		}
	} else {
		printf("upgrade kernel failed: nand or sata?\r\n");
		return -1;
	}

	printf("upgrade kernel success\r\n");

	return ret;
err:
	printf("upgrade kernel failed!\r\n");
	return ret;
}

static int update_rootfs(int dev, char *typename)
{
	int ret = -1;

	printf("update rootfs to nand............\n");

	switch (dev) {
	case UPDATE_DEV_USB:
		ret = run_command("fatload usb 0:1 ${loadaddr} update/rootfs-ubifs-ze.img", 0);
		break;
	case UPDATE_DEV_TFTP:
		ret = run_command("tftpboot ${loadaddr} rootfs-ubifs-ze.img", 0);
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("fatload mmc 0:1 ${loadaddr} update/rootfs-ubifs-ze.img", 0);
		break;
	}

	if (ret) {
		printf("upgrade rootfs failed no update/rootfs-ubifs-ze.img file!\r\n");
		return ret;
	}

	ret = mtdparts_probe();
	if (ret) {
		return ret;
	}

	ret = run_command("nand erase.part root;nand write ${loadaddr} root ${filesize}", 0);
	if (ret) {
		printf("upgrade rootfs failed!\r\n");
	} else {
		printf("upgrade rootfs success\r\n");
	}

	return ret;
}

static int do_loongson_update(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	int ret;
	int dev = UPDATE_DEV_UNKNOWN;
	int type = UPDATE_TYPE_UNKNOWN;

	for (dev = 0; dev < UPDATE_DEV_COUNT; dev++) {
		if (update_devname_str[dev] &&
			!strcmp(argv[1], update_devname_str[dev]))
			break;
	}
	if ((dev == UPDATE_DEV_UNKNOWN) || (dev >= UPDATE_DEV_COUNT)) {
		printf("upgrade dev unknown\r\n");
		printf("please intput dev: usb\\tftp\\mmc\r\n");
		goto err;
	}

	for (type = 0; type < UPDATE_TYPE_COUNT; type++) {
		if (update_typename_str[type] &&
			!strcmp(argv[2], update_typename_str[type]))
			break;
	}
	if ((type == UPDATE_TYPE_UNKNOWN) || (type >= UPDATE_TYPE_COUNT)) {
		printf("upgrade type unknown\r\n");
		printf("please intput type: kernel\\rootfs\\uboot\\dtb\\all\r\n");
		goto err;
	}

	switch (dev) {
	case UPDATE_DEV_USB:
		run_command("usb start", 0);
		ret = run_command("usb storage", 0);
		if (ret) {
			printf("upgrade failed no usb storage!\r\n");
			goto err;
		}
		break;
	case UPDATE_DEV_TFTP:
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("mmc rescan", 0);
		if (ret) {
			printf("upgrade failed no mmc storage!\r\n");
			goto err;
		}
		break;
	default:
		break;
	}

	switch (type) {
	case UPDATE_TYPE_UBOOT:
		ret = update_uboot(dev);
		break;
	case UPDATE_TYPE_DTB:
		ret = update_dtb(dev);
		break;
	case UPDATE_TYPE_KERNEL:
		ret = update_kernel(dev, argv[3], argv[4]);
		break;
	case UPDATE_TYPE_ROOTFS:
		ret = update_rootfs(dev, NULL);
		break;
	case UPDATE_TYPE_ALL:
		ret = update_uboot(dev);
		ret = update_dtb(dev);
		ret = update_kernel(dev, argv[3], argv[4]);
		ret = update_rootfs(dev, NULL);
		break;
	default:
		break;
	}

	return ret;

err:
	return -1;
}

U_BOOT_CMD(
	loongson_update,    5,    1,     do_loongson_update,
	"upgrade uboot kernel rootfs dtb over usb\\tftp\\mmc.",
	""
);
